package contest.p1530;

import ans.base.TreeNode;

import java.util.HashMap;
import java.util.Map;

class Solution {
    public int countPairs(TreeNode root, int distance) {

        int[] ans = new int[1];
        dfs(root, distance, ans);
        return ans[0];
    }

    /**
     * 返回深度到该节点的 子树个数 k,v 深度,个数
     * @param root
     * @param distance
     * @param ans
     * @return
     */
    public Map<Integer, Integer> dfs(TreeNode root, int distance, int[] ans) {
        if (root == null) {
            return new HashMap<>(0);
        }
        if (root.left == null && root.right == null) {
            Map<Integer, Integer> map = new HashMap<>(1);
            //深度,个数
            map.put(0, 1);
            return map;
        }
        //获取左子树深度和个数对应的MAP
        Map<Integer, Integer> dfs = dfs(root.left, distance, ans);
        //获取右子树深度和个数对应的MAP
        Map<Integer, Integer> dfs1 = dfs(root.right, distance, ans);
        //筛选出符合条件的记录到结果里 ans[0]
        for (Map.Entry<Integer, Integer> entry : dfs.entrySet()) {
            for (Map.Entry<Integer, Integer> entry1 : dfs1.entrySet()) {
                if (entry.getKey() + entry1.getKey() + 2 <= distance) {
                    ans[0] += entry.getValue() * entry1.getValue();
                }
            }
        }

        return combineResult(dfs, dfs1);
    }

    /**
     * 合并左右子树深度记录的MAP 并且高度+1
     *
     * @param dfs0
     * @param dfs1
     */
    private Map<Integer, Integer> combineResult(Map<Integer, Integer> dfs0, Map<Integer, Integer> dfs1) {
        Map<Integer, Integer> result = new HashMap<>();
        for (Map.Entry<Integer, Integer> entry : dfs0.entrySet()) {
            result.put(entry.getKey() + 1, entry.getValue());
        }
        for (Map.Entry<Integer, Integer> entry : dfs1.entrySet()) {
            Integer integer = result.get(entry.getKey() + 1);
            if (integer == null) {
                result.put(entry.getKey() + 1, entry.getValue());
            } else {
                result.put(entry.getKey() + 1, entry.getValue() + integer);
            }
        }
        return result;
    }
}

